import contextlib
import io
import operator
import os
import pathlib
import tarfile
import typing as tp

import docker
from loguru import logger


class DockerClient:
    IMAGE_PREFIX = "image-"

    def __init__(self, *args, **kwargs):
        self._proxy = docker.DockerClient(*args, **kwargs)

    @classmethod
    def from_env(cls):
        return cls(**dict(os.environ))

    def ensure_docker_image(self, fh, image_path=None) -> tp.Optional[pathlib.Path]:
        res = self._proxy.api.build(fileobj=fh, decode=True)

        image = None

        for s in res:
            if "stream" in s:
                for l in s["stream"].strip().splitlines():
                    logger.info(f"{l}")

            if "aux" in s and "ID" in s["aux"]:
                image = s["aux"]["ID"]

        if not image:
            raise Exception("unable to determine built Docker image")

        if image_path:
            tar_path = pathlib.Path(str(image_path) + ".tar")
            with tar_path.open("wb") as fh:
                for chunk in self._proxy.images.get(image).save():
                    fh.write(chunk)

            with image_path.open("w") as fh:
                fh.write(image + "\n")

        return image

    def build_docker_image(self, image_data: bytes, image_dir: pathlib.Path, name) -> tp.Optional[pathlib.Path]:
        image_path = image_dir / f"{self.IMAGE_PREFIX}{name}"

        return self.ensure_docker_image(io.BytesIO(image_data), image_path=image_path)

    def create_image(self, dockerfile: pathlib.Path, image_dir: pathlib.Path, name: str) -> tp.Optional[pathlib.Path]:
        with dockerfile.open("rb") as fh:
            image_data = fh.read()
        return self.build_docker_image(image_data, image_dir, name)
